package main

import (
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// AppRecord 表示应用记录
type AppRecord struct {
	ID               string    `bson:"id,omitempty"`
	TrackName        string    `bson:"track_name,omitempty"`
	SizeBytes        int64     `bson:"size_bytes,omitempty"`
	Currency         string    `bson:"currency,omitempty"`
	Price            float64   `bson:"price,omitempty"`
	RatingCountTotal int       `bson:"rating_count_total,omitempty"`
	RatingCountVer   int       `bson:"rating_count_ver,omitempty"`
	UserRating       float64   `bson:"user_rating,omitempty"`
	UserRatingVer    float64   `bson:"user_rating_ver,omitempty"`
	Version          string    `bson:"version,omitempty"`
	ContentRating    string    `bson:"content_rating,omitempty"`
	PrimeGenre       string    `bson:"prime_genre,omitempty"`
	SupportedDevices int       `bson:"supported_devices,omitempty"`
	IpadScreenShots  int       `bson:"ipad_screenshots,omitempty"`
	Languages        int       `bson:"languages,omitempty"`
	VPP_LICENSE      int       `bson:"vpp_license,omitempty"`
	CreatedAt        time.Time `bson:"created_at,omitempty"`
}

const (
	mongoURI       = "mongodb://admin:password@localhost:27017" // MongoDB连接URI
	databaseName   = "appstore"                  // 数据库名
	collectionName = "app_records"               // 集合名
	batchSize      = 100                         // 批量插入大小
	timeout        = 30 * time.Second            // 超时时间
)

func main() {
	if len(os.Args) != 2 {
		log.Fatal("用法: go run main.go <csv文件路径>")
	}

	csvFile := os.Args[1]

	// 读取CSV文件
	records, err := readCSVFile(csvFile)
	if err != nil {
		log.Fatalf("读取CSV文件失败: %v", err)
	}

	fmt.Printf("成功读取 %d 条记录\n", len(records))

	// 连接到MongoDB
	client, err := connectToMongoDB()
	if err != nil {
		log.Fatalf("连接MongoDB失败: %v", err)
	}
	defer client.Disconnect(context.Background())

	fmt.Println("成功连接到MongoDB")

	// 获取集合
	collection := client.Database(databaseName).Collection(collectionName)

	// 上传数据到MongoDB
	if err := uploadToMongoDB(collection, records); err != nil {
		log.Fatalf("上传数据到MongoDB失败: %v", err)
	}

	fmt.Printf("成功上传 %d 条记录到MongoDB集合 '%s.%s'\n", len(records), databaseName, collectionName)
}

// connectToMongoDB 连接到MongoDB
func connectToMongoDB() (*mongo.Client, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	clientOptions := options.Client().ApplyURI(mongoURI)

	// 设置连接选项
	clientOptions.SetMaxPoolSize(100)
	clientOptions.SetMinPoolSize(10)
	clientOptions.SetMaxConnIdleTime(5 * time.Minute)

	client, err := mongo.Connect(ctx, clientOptions)
	if err != nil {
		return nil, err
	}

	// 检查连接
	err = client.Ping(ctx, nil)
	if err != nil {
		return nil, err
	}

	return client, nil
}

// readCSVFile 读取CSV文件并解析为AppRecord切片
func readCSVFile(filename string) ([]AppRecord, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	reader := csv.NewReader(file)

	// 读取表头（跳过）
	headers, err := reader.Read()
	if err != nil {
		return nil, err
	}

	fmt.Printf("CSV表头: %v\n", headers)

	var records []AppRecord
	lineNumber := 1 // 从第2行开始（跳过表头）

	for {
		line, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Printf("读取第%d行失败: %v, 跳过该行", lineNumber, err)
			lineNumber++
			continue
		}

		// 解析CSV行到AppRecord结构体
		record, err := parseCSVLine(line)
		if err != nil {
			log.Printf("解析第%d行失败: %v, 跳过该行", lineNumber, err)
			lineNumber++
			continue
		}

		records = append(records, record)
		lineNumber++
	}

	return records, nil
}

// parseCSVLine 解析CSV行到AppRecord结构体
func parseCSVLine(line []string) (AppRecord, error) {
	var record AppRecord
	var err error

	if len(line) < 16 {
		return record, fmt.Errorf("行数据不完整，期望16列，实际%d列", len(line))
	}

	record.ID = strings.TrimSpace(line[0])
	record.TrackName = strings.TrimSpace(line[1])

	// 解析SizeBytes
	if line[2] != "" {
		record.SizeBytes, err = strconv.ParseInt(line[2], 10, 64)
		if err != nil {
			return record, fmt.Errorf("解析SizeBytes失败: %v", err)
		}
	}

	record.Currency = strings.TrimSpace(line[3])

	// 解析Price
	if line[4] != "" {
		record.Price, err = strconv.ParseFloat(line[4], 64)
		if err != nil {
			return record, fmt.Errorf("解析Price失败: %v", err)
		}
	}

	// 解析RatingCountTotal
	if line[5] != "" {
		record.RatingCountTotal, err = strconv.Atoi(line[5])
		if err != nil {
			return record, fmt.Errorf("解析RatingCountTotal失败: %v", err)
		}
	}

	// 解析RatingCountVer
	if line[6] != "" {
		record.RatingCountVer, err = strconv.Atoi(line[6])
		if err != nil {
			return record, fmt.Errorf("解析RatingCountVer失败: %v", err)
		}
	}

	// 解析UserRating
	if line[7] != "" {
		record.UserRating, err = strconv.ParseFloat(line[7], 64)
		if err != nil {
			return record, fmt.Errorf("解析UserRating失败: %v", err)
		}
	}

	// 解析UserRatingVer
	if line[8] != "" {
		record.UserRatingVer, err = strconv.ParseFloat(line[8], 64)
		if err != nil {
			return record, fmt.Errorf("解析UserRatingVer失败: %v", err)
		}
	}

	record.Version = strings.TrimSpace(line[9])
	record.ContentRating = strings.TrimSpace(line[10])
	record.PrimeGenre = strings.TrimSpace(line[11])

	// 解析SupportedDevices
	if line[12] != "" {
		record.SupportedDevices, err = strconv.Atoi(line[12])
		if err != nil {
			return record, fmt.Errorf("解析SupportedDevices失败: %v", err)
		}
	}

	// 解析IpadScreenShots
	if line[13] != "" {
		record.IpadScreenShots, err = strconv.Atoi(line[13])
		if err != nil {
			return record, fmt.Errorf("解析IpadScreenShots失败: %v", err)
		}
	}

	// 解析Languages
	if line[14] != "" {
		record.Languages, err = strconv.Atoi(line[14])
		if err != nil {
			return record, fmt.Errorf("解析Languages失败: %v", err)
		}
	}

	// 解析VPP_LICENSE
	if line[15] != "" {
		record.VPP_LICENSE, err = strconv.Atoi(line[15])
		if err != nil {
			return record, fmt.Errorf("解析VPP_LICENSE失败: %v", err)
		}
	}

	// 设置创建时间
	record.CreatedAt = time.Now()

	return record, nil
}

// uploadToMongoDB 批量上传记录到MongoDB
func uploadToMongoDB(collection *mongo.Collection, records []AppRecord) error {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// 批量插入
	var documents []interface{}
	totalRecords := len(records)

	for i, record := range records {
		documents = append(documents, record)

		// 达到批量大小或最后一条记录时执行插入
		if len(documents) >= batchSize || i == totalRecords-1 {
			result, err := collection.InsertMany(ctx, documents)
			if err != nil {
				return fmt.Errorf("批量插入失败: %v", err)
			}

			fmt.Printf("已插入 %d 条记录，进度: %d/%d (%.1f%%)\n",
				len(result.InsertedIDs), i+1, totalRecords, float64(i+1)/float64(totalRecords)*100)

			// 清空当前批次
			documents = nil
		}
	}

	return nil
}

// 可选：创建索引的函数
func createIndexes(collection *mongo.Collection) error {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// 创建索引
	indexes := []mongo.IndexModel{
		{
			Keys:    bson.D{{Key: "id", Value: 1}},
			Options: options.Index().SetUnique(true).SetName("id_unique"),
		},
		{
			Keys:    bson.D{{Key: "track_name", Value: 1}},
			Options: options.Index().SetName("track_name_idx"),
		},
		{
			Keys:    bson.D{{Key: "prime_genre", Value: 1}},
			Options: options.Index().SetName("prime_genre_idx"),
		},
		{
			Keys:    bson.D{{Key: "user_rating", Value: -1}},
			Options: options.Index().SetName("user_rating_desc_idx"),
		},
		{
			Keys:    bson.D{{Key: "price", Value: 1}},
			Options: options.Index().SetName("price_idx"),
		},
	}

	_, err := collection.Indexes().CreateMany(ctx, indexes)
	if err != nil {
		return fmt.Errorf("创建索引失败: %v", err)
	}

	fmt.Println("索引创建成功")
	return nil
}
